# coding=utf-8
from math import exp
from numpy import *
import matplotlib.pyplot as plt


# create test data.
def test_data():
    data_mat = []
    label_mat = []

    data_file = open("data/testSet.txt")
    for line in data_file.readlines():
        line_arr = line.strip().split()
        data_mat.append([1.0, float(line_arr[0]), float(line_arr[1])])
        label_mat.append(int(line_arr[2]))
    return data_mat, label_mat


# sigmoid function.
def sigmoid(x):
    return 1.0 / (1 + exp(-x))


# use gradient ascent method to get the best weights.
def grad_ascent(data_set, labels_set):
    data_mat = mat(data_set)
    label_mat = mat(labels_set).transpose()

    m, n = shape(data_mat)

    alpha = 0.001
    max_cycles = 500  # the maximum number of iterations.
    weights = ones((n, 1))

    for k in range(max_cycles):
        h = sigmoid(data_mat * weights)  # current predictive values.
        error = (label_mat - h)
        # Adjust the weight in the direction of the error.
        weights = weights + alpha * data_mat.transpose() * error

    return weights


def plot_best(weights):
    data_mat, label_mat = test_data()

    data_arr = array(data_mat)
    n = shape(data_arr)[0]

    xcord1 = []
    ycord1 = []

    xcord2 = []
    ycord2 = []

    for i in range(n):
        if int(label_mat[i]) == 1:
            xcord1.append(data_arr[i, 1])
            ycord1.append(data_arr[i, 2])
        else:
            xcord2.append(data_arr[i, 1])
            ycord2.append(data_arr[i, 2])

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(xcord1, ycord1, s=30, c='red', marker='s')
    ax.scatter(xcord2, ycord2, s=30, c='green')

    x = arange(-3.0, 3.0, 0.1)
    y = (-weights[0] - weights[1] * x) / weights[2]

    ax.plot(x, y)
    plt.xlabel('X1')
    plt.ylabel('X2')

    plt.show()


data, labels = test_data()
weights = grad_ascent(data, labels)
plot_best(weights.getA())

